/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <filesystem>
#include <memory>
#include <set>
#include <string>
#include <boost/algorithm/string.hpp>
#include <boost/algorithm/string/predicate.hpp>
#include <boost/algorithm/string/replace.hpp>
#include <fmt/format.h>
#include <thrift/compiler/ast/t_struct.h>
#include <thrift/compiler/generate/go/util.h>
#include <thrift/compiler/generate/t_mstch_generator.h>

namespace apache::thrift::compiler {

namespace {

class t_mstch_go_generator : public t_mstch_generator {
 public:
  using t_mstch_generator::t_mstch_generator;

  whisker_options render_options() const override {
    whisker_options opts;
    opts.allowed_undefined_variables = {
        "service:autogen_path",
    };
    return opts;
  }

  std::string template_prefix() const override { return "go"; }

  void generate_program() override;

 private:
  void initialize_context(context_visitor& visitor) override {
    data_.set_current_program(program_);
    data_.compute_go_package_aliases();
    data_.compute_struct_to_field_names();
    data_.compute_req_resp_structs();
    data_.compute_thrift_metadata_types();

    if (auto gen_metadata = get_option("gen_metadata")) {
      data_.gen_metadata = (gen_metadata.value() == "true");
    }
    if (auto gen_default_get = get_option("gen_default_get")) {
      data_.gen_default_get = (gen_default_get.value() == "true");
    }
    if (auto use_reflect_codec = get_option("use_reflect_codec")) {
      data_.use_reflect_codec = (use_reflect_codec.value() == "true");
    }

    // Visit the transient request/response structs generated by this generator
    for (const t_struct* s : data_.req_resp_structs) {
      whisker_generator_visitor_context ctx;
      visitor(ctx, *s);
    }
  }

  go::codegen_data data_;

  // Whisker prototype helpers
  std::string go_qualified_name_(const t_named& self) const {
    auto prefix = data_.go_package_alias_prefix(self.program());
    return prefix + go::go_name(self);
  }
  bool is_field_inside_union(const t_field& field) const {
    // Whether field is part of a union
    const t_structured* parent = context().get_field_parent(&field);
    assert(parent != nullptr);
    return parent->is<t_union>();
  }
  bool is_field_pointer(const t_field& field) const {
    // Whether this field is a pointer '*' in a Go struct definition:
    //  * Struct-type fields are pointers.
    //     * Union-type fields are pointers too - by extension.
    //  * Fields inside a union are pointers.
    //     * Except (!!!) when the underlying type itself is nilable (map/slice)
    //  * Optional fields are pointers.
    //     * Except (!!!) when the underlying type itself is nilable (map/slice)
    auto real_type = field.type()->get_true_type();
    return go::is_type_go_struct(real_type) ||
        ((is_field_inside_union(field) ||
          field.qualifier() == t_field_qualifier::optional) &&
         !go::is_type_go_nilable(real_type));
  }

  prototype<t_program>::ptr make_prototype_for_program(
      const prototype_database& proto) const override {
    auto base = t_whisker_generator::make_prototype_for_program(proto);
    auto def = whisker::dsl::prototype_builder<h_program>::extends(base);

    def.property("go_pkg_name", [this](const t_program& self) {
      auto pkg_name = go::get_go_package_base_name(&self);
      if (data_.compat) {
        return pkg_name;
      } else {
        return go::snakecase(pkg_name);
      }
    });
    def.property("go_package_alias", [this](const t_program& self) {
      return data_.get_go_package_alias(&self);
    });
    def.property("thirft_source_path", [](const t_program& self) {
      return self.path();
    });
    def.property("compat?", [this](const t_program&) { return data_.compat; });
    def.property("compat_setters?", [this](const t_program&) {
      return data_.compat_setters;
    });
    def.property("gen_metadata?", [this](const t_program&) {
      return data_.gen_metadata;
    });
    def.property("gen_default_get?", [this](const t_program&) {
      return data_.gen_default_get;
    });

    return std::move(def).make();
  }

  prototype<t_interface>::ptr make_prototype_for_interface(
      const prototype_database& proto) const override {
    auto base = t_whisker_generator::make_prototype_for_interface(proto);
    auto def = whisker::dsl::prototype_builder<h_interface>::extends(base);

    def.property("go_name", [](const t_interface& self) {
      return go::munge_ident(self.name());
    });
    def.property("go_client_name", [](const t_interface& self) {
      return go::munge_ident(self.name(), /* exported = */ false) +
          "ClientImpl";
    });

    return std::move(def).make();
  }

  prototype<t_function>::ptr make_prototype_for_function(
      const prototype_database& proto) const override {
    auto base = t_whisker_generator::make_prototype_for_function(proto);
    auto def = whisker::dsl::prototype_builder<h_function>::extends(base);

    def.property("go_name", [](const t_function& self) {
      return go::get_go_func_name(&self);
    });
    def.property("go_client_supported?", [](const t_function& self) {
      return go::is_func_go_client_supported(&self);
    });
    def.property("go_server_supported?", [](const t_function& self) {
      return go::is_func_go_server_supported(&self);
    });
    def.property("retval_field_name", [](const t_function&) {
      // Field name for the return value.
      return go::munge_ident(go::DEFAULT_RETVAL_FIELD_NAME, /*exported*/ true);
    });
    def.property("ctx_arg_name", [](const t_function& self) {
      // This helper returns the Context object name to be used in the function
      // signature. "ctx" by default, "ctx<num>" in case of name collisions with
      // other function arguments. The name is guaranteed not to collide.
      return go::get_go_func_unique_arg_name(&self, "ctx");
    });

    return std::move(def).make();
  }

  prototype<t_type>::ptr make_prototype_for_type(
      const prototype_database& proto) const override {
    auto base = t_whisker_generator::make_prototype_for_type(proto);
    auto def = whisker::dsl::prototype_builder<h_type>::extends(base);

    def.property("nilable?", [](const t_type& self) {
      auto real_type = self.get_true_type();
      return go::is_type_go_nilable(real_type);
    });
    def.property("go_comparable?", [](const t_type& self) {
      return go::is_type_go_comparable(&self);
    });
    def.property(
        "named?", [](const t_type& self) { return !self.name().empty(); });
    def.property("metadata_primitive?", [](const t_type& self) {
      // Whether this type is primitive from metadata.thrift perspective.
      // i.e. see ThriftPrimitiveType enum in metadata.thrift
      auto real_type = self.get_true_type();
      return go::is_type_metadata_primitive(real_type);
    });
    def.property("metadata_name", [](const t_type& self) {
      return go::get_go_type_metadata_name(self);
    });
    def.property("codec_type_spec_name", [](const t_type& self) {
      return go::get_go_type_codec_type_spec_name(self);
    });
    def.property("metadata_thrift_type_getter", [this](const t_type& self) {
      // Program will be null for primitive (base) types.
      // They should be treated as being from the current program.
      auto is_from_current_program =
          self.program() == nullptr || data_.is_current_program(self.program());

      if (is_from_current_program) {
        // If the type is from the current program, we can simply use its
        // corresponding *ThriftType variable already present in the program.
        return go::get_go_type_metadata_name(self);
      } else {
        // If the type is external, we must retrieve it from its corresponding
        // program/package using GetMetadataThriftType helper method.
        return fmt::format(
            "{}.GetMetadataThriftType(\"{}\")",
            data_.get_go_package_alias(self.program()),
            self.get_full_name());
      }
    });
    def.property("codec_type_spec_getter", [this](const t_type& self) {
      // Program will be null for primitive (base) types.
      // They should be treated as being from the current program.
      auto is_from_current_program =
          self.program() == nullptr || data_.is_current_program(self.program());

      if (is_from_current_program) {
        // If the type is from the current program, we can simply use its
        // corresponding *ThriftType variable already present in the program.
        return go::get_go_type_codec_type_spec_name(self);
      } else {
        // If the type is external, we must retrieve it from its corresponding
        // program/package using GetMetadataThriftType helper method.
        return fmt::format(
            "{}.GetCodecTypeSpec(\"{}\")",
            data_.get_go_package_alias(self.program()),
            self.get_full_name());
      }
    });

    return std::move(def).make();
  }

  prototype<t_typedef>::ptr make_prototype_for_typedef(
      const prototype_database& proto) const override {
    auto base = t_whisker_generator::make_prototype_for_typedef(proto);
    auto def = whisker::dsl::prototype_builder<h_typedef>::extends(base);

    def.property("defined_kind?", [](const t_typedef& self) {
      // NOTE: there are multiple typedef "kinds":
      //  * defined - typedef actually defined in a Thrift schema by a human.
      //  * unnamed - typedef used for unstructured annotations.
      //  * placeholder - typedef used as a placeholder during AST parsing
      //    when not all type are fully known yet. During generation, when we
      //    encounter this kind fo typedef, we should skip it to the underlying
      //    "real" type or "defined" typedef to ensure code correctness.
      return self.typedef_kind() == t_typedef::kind::defined;
    });
    def.property("go_qualified_new_func", [this](const t_typedef& self) {
      auto prefix = data_.go_package_alias_prefix(self.program());
      return fmt::format("{}New{}", prefix, go::go_name(self));
    });
    def.property("go_qualified_write_func", [this](const t_typedef& self) {
      auto prefix = data_.go_package_alias_prefix(self.program());
      return fmt::format("{}Write{}", prefix, go::go_name(self));
    });
    def.property("go_qualified_read_func", [this](const t_typedef& self) {
      auto prefix = data_.go_package_alias_prefix(self.program());
      return fmt::format("{}Read{}", prefix, go::go_name(self));
    });

    return std::move(def).make();
  }

  prototype<t_structured>::ptr make_prototype_for_structured(
      const prototype_database& proto) const override {
    auto base = t_whisker_generator::make_prototype_for_structured(proto);
    auto def = whisker::dsl::prototype_builder<h_structured>::extends(base);

    def.property("go_public_req_name", [](const t_structured& self) {
      return boost::algorithm::erase_first_copy(self.name(), "req") +
          "ArgsDeprecated";
    });
    def.property("go_public_resp_name", [](const t_structured& self) {
      return boost::algorithm::erase_first_copy(self.name(), "resp") +
          "ResultDeprecated";
    });
    def.property("struct_spec_name", [](const t_structured& self) {
      return fmt::format("premadeStructSpec_{}", self.name());
    });
    def.property("struct_metadata_name", [](const t_structured& self) {
      return fmt::format("premadeStructMetadata_{}", self.name());
    });
    def.property("use_reflect_codec?", [this](const t_structured& self) {
      auto use_reflect_codec_annotation =
          self.find_structured_annotation_or_null(kGoUseReflectCodecUri);
      return data_.use_reflect_codec || use_reflect_codec_annotation != nullptr;
    });
    def.property("go_qualified_new_func", [this](const t_structured& self) {
      return fmt::format(
          "{}{}{}",
          data_.go_package_alias_prefix(self.program()),
          data_.is_req_resp_struct(self) ? "new" : "New",
          go::munge_ident(self.name(), /*exported=*/true));
    });

    return std::move(def).make();
  }

  prototype<t_field>::ptr make_prototype_for_field(
      const prototype_database& proto) const override {
    auto base = t_whisker_generator::make_prototype_for_field(proto);
    auto def = whisker::dsl::prototype_builder<h_field>::extends(base);

    def.property("go_name", [](const t_field& self) {
      return go::get_go_field_name(&self);
    });
    def.property("go_arg_name", [](const t_field& self) {
      auto arg_name = go::munge_ident(self.name(), /*exported*/ false);
      // Avoid 'context' package import collision
      if (arg_name == "context") {
        arg_name += "_";
      }
      return arg_name;
    });
    def.property("go_setter_name", [this](const t_field& self) {
      auto setter_name = "Set" + go::get_go_field_name(&self);
      // Setters which collide with existing field names should be suffixed with
      // an underscore.
      const t_structured* parent = context().get_field_parent(&self);
      assert(parent != nullptr);
      if (auto stfn_iter = data_.struct_to_field_names.find(parent);
          stfn_iter != data_.struct_to_field_names.end()) {
        while (stfn_iter->second.count(setter_name) > 0) {
          setter_name += "_";
        }
      }

      return setter_name;
    });
    def.property("key_str", [](const t_field& self) {
      // Legacy schemas may have negative tags - replace minus with an
      // underscore.
      if (self.id() < 0) {
        return "_" + std::to_string(-self.id());
      } else {
        return std::to_string(self.id());
      }
    });
    def.property("retval?", [](const t_field& self) {
      return self.name() == go::DEFAULT_RETVAL_FIELD_NAME;
    });
    def.property("go_tag?", [](const t_field& self) {
      return go::get_go_tag_annotation(&self) != nullptr;
    });
    def.property("go_tag", [](const t_field& self) {
      auto tag = go::get_go_tag_annotation(&self);
      if (tag != nullptr) {
        return *tag;
      }
      return std::string();
    });
    def.property("nilable?", [this](const t_field& self) {
      // Whether this field can be set to 'nil' in Go:
      //  * Fields of nilable Go types can be set to 'nil' (map/slice/struct)
      //  * Fields inside a union can be set to 'nil' ('is_pointer' above)
      //  * Optional fields can be set to 'nil' (see 'is_pointer' above)
      auto real_type = self.type()->get_true_type();
      return go::is_type_go_nilable(real_type) || is_field_inside_union(self) ||
          self.qualifier() == t_field_qualifier::optional;
    });
    def.property("must_be_set_to_serialize?", [this](const t_field& self) {
      // Whether the field must be set (non-nil) in order to serialize:
      //  * Struct type fields must be set (to avoid nil pointer dereference)
      //  * Fields inside a union must be set (that's the point of a union)
      //  * Optional fields must be set ("unset" optional fields must not be
      //  serailized as per Thrift-spec)
      auto real_type = self.type()->get_true_type();
      return go::is_type_go_struct(real_type) || is_field_inside_union(self) ||
          self.qualifier() == t_field_qualifier::optional;
    });
    def.property("pointer?", [this](const t_field& self) {
      // See comment in the helper method for details.
      return is_field_pointer(self);
    });
    def.property("non_struct_pointer?", [this](const t_field& self) {
      // Whether this field is a non-struct pointer.
      auto real_type = self.type()->get_true_type();
      return is_field_pointer(self) && !go::is_type_go_struct(real_type);
    });
    def.property("json_omitempty?", [this](const t_field& self) {
      // Whether this field should be tagged with 'json:"omitempty"'.
      // Optional and union fields should be tagged as such.
      return self.qualifier() == t_field_qualifier::optional ||
          is_field_inside_union(self);
    });

    return std::move(def).make();
  }

  prototype<t_const>::ptr make_prototype_for_const(
      const prototype_database& proto) const override {
    auto base = t_mstch_generator::make_prototype_for_const(proto);
    auto def = whisker::dsl::prototype_builder<h_const>::extends(base);

    def.property("var?", [](const t_const& self) {
      // go_var returns true to use a var instead of a const in Go for the
      // thrift const definition (e.g for structs, maps, or lists which
      // cannot be const in go)
      auto real_type = self.type()->get_true_type();
      return go::is_type_go_nilable(real_type);
    });

    return std::move(def).make();
  }

  prototype<t_const_value>::ptr make_prototype_for_const_value(
      const prototype_database& proto) const override {
    auto base = t_mstch_generator::make_prototype_for_const_value(proto);
    auto def = whisker::dsl::prototype_builder<h_const_value>::extends(base);

    def.property("go_quoted_value", [](const t_const_value& self) {
      return go::quote(self.get_string());
    });

    return std::move(def).make();
  }

  prototype<t_named>::ptr make_prototype_for_named(
      const prototype_database& proto) const override {
    auto base = t_whisker_generator::make_prototype_for_named(proto);
    auto def = whisker::dsl::prototype_builder<h_named>::extends(base);

    def.property(
        "go_name", [](const t_named& self) { return go::go_name(self); });
    def.property("go_qualified_name", [this](const t_named& self) {
      return go_qualified_name_(self);
    });
    def.property("go_package_alias_prefix", [this](const t_named& self) {
      return data_.go_package_alias_prefix(self.program());
    });
    def.property(
        "docs", [](const t_named& self) { return go::doc_comment(&self); });

    return std::move(def).make();
  }
};

class mstch_go_program : public mstch_program {
 public:
  mstch_go_program(
      const t_program* p,
      mstch_context& ctx,
      mstch_element_position pos,
      go::codegen_data* data)
      : mstch_program(p, ctx, pos), data_(*data) {
    register_methods(
        this,
        {
            {"program:go_import_path", &mstch_go_program::go_import_path},
            {"program:thrift_imports", &mstch_go_program::thrift_imports},
            {"program:import_metadata_package?",
             &mstch_go_program::should_import_metadata_package},
            {"program:metadata_qualifier",
             &mstch_go_program::metadata_qualifier},
            {"program:thrift_metadata_types",
             &mstch_go_program::thrift_metadata_types},
            {"program:req_resp_structs", &mstch_go_program::req_resp_structs},
        });
  }

  mstch::node thrift_imports() {
    mstch::array a;
    for (const auto* program : program_->get_includes_for_codegen()) {
      a.emplace_back(make_mstch_program_cached(program, context_));
    }
    return a;
  }
  mstch::node go_import_path() { return get_go_import_path_(); }
  mstch::node should_import_metadata_package() {
    // We don't need to import the metadata package if we are
    // generating metadata inside the metadata package itself. Duh.
    return !is_metadata_package_();
  }
  mstch::node metadata_qualifier() {
    // We don't need to use "metadata." qualifier when generating
    // metadata inside the metadata package itself.
    if (!is_metadata_package_()) {
      return std::string("metadata.");
    } else {
      return std::string("");
    }
  }
  mstch::node thrift_metadata_types() {
    return make_mstch_array(
        data_.thrift_metadata_types, *context_.type_factory);
  }
  mstch::node req_resp_structs() {
    return make_mstch_array(data_.req_resp_structs, *context_.struct_factory);
  }

 private:
  go::codegen_data& data_;

  std::string get_go_import_path_() { return go::get_go_package_dir(program_); }
  bool is_metadata_package_() {
    return get_go_import_path_() == go::THRIFT_METADATA_IMPORT;
  }
};

class mstch_go_const_value : public mstch_const_value {
 public:
  mstch_go_const_value(
      const t_const_value* cv,
      mstch_context& ctx,
      mstch_element_position pos,
      const t_const* current_const,
      const t_type* expected_type)
      : mstch_const_value(cv, ctx, pos, current_const, expected_type) {}

  bool same_type_as_expected() const override { return true; }
};

class mstch_go_struct : public mstch_struct {
 public:
  mstch_go_struct(
      const t_structured* s,
      mstch_context& ctx,
      mstch_element_position pos,
      go::codegen_data* data)
      : mstch_struct(s, ctx, pos), data_(*data) {
    register_methods(
        this,
        {
            {"struct:go_name", &mstch_go_struct::go_name},
            {"struct:go_qualified_name", &mstch_go_struct::go_qualified_name},
            {"struct:req_resp?", &mstch_go_struct::is_req_resp_struct},
            {"struct:resp?", &mstch_go_struct::is_resp_struct},
            {"struct:req?", &mstch_go_struct::is_req_struct},
            {"struct:stream?", &mstch_go_struct::is_stream_struct},
            {"struct:fields_sorted", &mstch_go_struct::fields_sorted},
        });
  }

  mstch::node go_name() { return go_name_(); }
  mstch::node go_qualified_name() {
    auto prefix = data_.go_package_alias_prefix(struct_->program());
    return prefix + go_name_();
  }
  mstch::node is_req_resp_struct() {
    // Whether this is a helper request or response struct.
    return data_.is_req_resp_struct(*struct_);
  }
  mstch::node is_resp_struct() {
    // Whether this is a helper response struct.
    return data_.is_req_resp_struct(*struct_) &&
        boost::algorithm::starts_with(struct_->name(), "resp");
  }
  mstch::node is_req_struct() {
    // Whether this is a helper request struct.
    return data_.is_req_resp_struct(*struct_) &&
        boost::algorithm::starts_with(struct_->name(), "req");
  }
  mstch::node is_stream_struct() {
    // Whether this is a helper stream struct.
    return data_.is_req_resp_struct(*struct_) &&
        boost::algorithm::starts_with(struct_->name(), "stream");
  }
  mstch::node fields_sorted() {
    // Fields (optionally) in the most optimal (memory-saving) layout order.
    if (struct_->has_structured_annotation(kGoMinimizePaddingUri)) {
      std::vector<const t_field*> fields_in_layout_order =
          struct_->fields_id_order();
      go::optimize_fields_layout(
          fields_in_layout_order, struct_->is<t_union>());
      return make_mstch_fields(fields_in_layout_order);
    }
    return make_mstch_fields(struct_->fields_id_order());
  }

 private:
  go::codegen_data& data_;

  std::string go_name_() {
    auto name = struct_->name();
    if (data_.is_req_resp_struct(*struct_)) {
      // Unexported/lowercase
      return go::munge_ident(name, false);
    } else {
      auto name_override = go::get_go_name_annotation(struct_);
      if (name_override != nullptr) {
        return *name_override;
      }
      // Exported/uppercase
      return go::munge_ident(name, true);
    }
  }
};

void t_mstch_go_generator::generate_program() {
  out_dir_base_ = "gen-go";

  mstch_context_.add<mstch_go_program>(&data_);
  mstch_context_.add<mstch_go_struct>(&data_);
  mstch_context_.add<mstch_go_const_value>();

  const auto& prog = cached_program(program_);
  render_to_file(prog, "const.go", "const.go");
  render_to_file(prog, "types.go", "types.go");
  render_to_file(prog, "svcs.go", "svcs.go");
  render_to_file(prog, "codec.go", "codec.go");
  if (data_.gen_metadata) {
    render_to_file(prog, "metadata.go", "metadata.go");
  }
}

} // namespace

THRIFT_REGISTER_GENERATOR(mstch_go, "Go", "");

} // namespace apache::thrift::compiler
